package io.kumao.gateway.filter;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.List;

import static org.springframework.http.HttpHeaders.AUTHORIZATION;
import static org.springframework.http.HttpStatus.UNAUTHORIZED;

/**
 * 令牌头全局过滤器.
 *
 * @author 崔作利
 */
@Slf4j
@Component
public class TokenHeaderGlobalFilter implements GlobalFilter, Ordered {

    private final static String DOMAIN = "domain";
    private final static String ACCESS_TOKEN_KEY_FORMAT = "access_token_%s";
    private final static String UNAUTHORIZED_JSON = "{\"success\":false,\"errorCode\":400,\"errorMessage\":\"unauthorized\",\"payload\":null}";

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest req = exchange.getRequest();
        ServerHttpResponse res = exchange.getResponse();

        ServerHttpRequest httpReq = req.mutate()
                .headers(headersConsumer -> headersConsumer.add(DOMAIN, req.getURI().getHost()))
                .build();

        ServerWebExchange webExchange = exchange.mutate().request(httpReq).build();

        HttpHeaders headers = req.getHeaders();
        List<String> auths = headers.get(AUTHORIZATION);
        if (CollectionUtils.isNotEmpty(auths)) {
            String accessTokenKey = String.format(ACCESS_TOKEN_KEY_FORMAT, auths.get(0));
            Boolean exists = stringRedisTemplate.hasKey(accessTokenKey);
            if (exists == null || !exists) {
                DataBufferFactory bufferFactory = res.bufferFactory();
                ServerHttpResponseDecorator httpRes = new ServerHttpResponseDecorator(res) {
                    @Override
                    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                        Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                        return super.writeWith(fluxBody.map(dataBuffer -> {
                            DataBufferUtils.release(dataBuffer);
                            return bufferFactory.wrap(UNAUTHORIZED_JSON.getBytes(StandardCharsets.UTF_8));
                        }));
                    }
                };
                httpRes.setStatusCode(UNAUTHORIZED);
                webExchange = exchange.mutate().request(httpReq).response(httpRes).build();
            }
        }
        return chain.filter(webExchange);
    }

    @Override
    public int getOrder() {
        return -2;
    }
}
